package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"strconv"
	"time"
)

const (
	CmdConnect uint8 = 1
	CmdBind    uint8 = 2
	CmdUdp     uint8 = 3
)
const (
	// CmdMuxBind is an extended SOCKS5 request CMD for
	// multiplexing transport with the binding server.
	CmdMuxBind uint8 = 0xF2
	// CmdUDPTun is an extended SOCKS5 request CMD for UDP over TCP.
	CmdUDPTun uint8 = 0xF3
)

const (
	AddrIPv4   uint8 = 1
	AddrDomain uint8 = 3
	AddrIPv6   uint8 = 4
)

var (
	ErrBadVersion  = errors.New("bad version")
	ErrBadFormat   = errors.New("bad format")
	ErrBadAddrType = errors.New("bad address type")
	ErrShortBuffer = errors.New("short buffer")
	ErrBadMethod   = errors.New("bad method")
	ErrAuthFailure = errors.New("auth failure")
)

//
// class def Addr begin
//
/*
Address
 +------+----------+----------+
 | ATYP |   ADDR   |   PORT   |
 +------+----------+----------+
 |  1   | Variable |    2     |
 +------+----------+----------+
*/
type Addr struct {
	Type uint8
	Host string
	Port uint16
}

func NewAddr(sa string) (addr *Addr, err error) {
	addr = &Addr{}
	err = addr.ParseFrom(sa)
	return
}

func (addr *Addr) ParseFrom(saddr string) error {
	host, sport, err := net.SplitHostPort(saddr)
	if err != nil {
		return err
	}
	port, err := strconv.Atoi(sport)
	if err != nil {
		return err
	}

	addr.Host = host
	addr.Port = uint16(port)
	addr.checkType()

	return nil
}

func (addr *Addr) ReadFrom(r io.Reader) (n int64, err error) {
	var b [255]byte

	_, err = io.ReadFull(r, b[:1])
	if err != nil {
		return
	}
	addr.Type = b[0]
	n++

	switch addr.Type {
	case AddrIPv4:
		_, err = io.ReadFull(r, b[:net.IPv4len])
		addr.Host = net.IP(b[:net.IPv4len]).String()
		n += net.IPv4len
	case AddrIPv6:
		_, err = io.ReadFull(r, b[:net.IPv6len])
		addr.Host = net.IP(b[:net.IPv6len]).String()
		n += net.IPv6len
	case AddrDomain:
		if _, err = io.ReadFull(r, b[:1]); err != nil {
			return
		}
		addrlen := int(b[0])
		n++

		_, err = io.ReadFull(r, b[:addrlen])
		addr.Host = string(b[:addrlen])
		n += int64(addrlen)
	default:
		err = ErrBadAddrType
		return
	}
	if err != nil {
		return
	}

	_, err = io.ReadFull(r, b[:2])
	addr.Port = binary.BigEndian.Uint16(b[:2])
	n += 2

	return
}

func (addr *Addr) WriteTo(w io.Writer) (int64, error) {
	var b [259]byte
	nn, err := addr.Encode(b[:])
	if err != nil {
		return int64(nn), err
	}

	nn, err = w.Write(b[:nn])
	return int64(nn), err
}

func (addr *Addr) Decode(b []byte) error {
	_, err := addr.ReadFrom(bytes.NewReader(b))
	return err
}

func (addr *Addr) Encode(b []byte) (int, error) {
	addr.checkType()

	b[0] = addr.Type
	pos := 1
	switch addr.Type {
	case AddrIPv4:
		ip4 := net.ParseIP(addr.Host).To4()
		if ip4 == nil {
			ip4 = net.IPv4zero.To4()
		}
		pos += copy(b[pos:], ip4)
	case AddrIPv6:
		ip16 := net.ParseIP(addr.Host).To16()
		if ip16 == nil {
			ip16 = net.IPv6zero.To16()
		}
		pos += copy(b[pos:], ip16)
	case AddrDomain:
		b[pos] = byte(len(addr.Host))
		pos++
		pos += copy(b[pos:], []byte(addr.Host))
	default:
		b[0] = AddrIPv4
		copy(b[pos:pos+net.IPv4len], net.IPv4zero.To4())
		pos += net.IPv4len
	}
	binary.BigEndian.PutUint16(b[pos:], addr.Port)
	pos += 2

	return pos, nil
}

func (addr *Addr) checkType() {
	switch addr.Type {
	case AddrIPv4, AddrIPv6, AddrDomain:
	default:
		addr.Type = AddrDomain
		if ip := net.ParseIP(addr.Host); ip != nil {
			if ip.To4() != nil {
				addr.Type = AddrIPv4
			} else {
				addr.Type = AddrIPv6
			}
		}
	}
}

func (addr *Addr) Length() (n int) {
	addr.checkType()

	switch addr.Type {
	case AddrIPv4:
		n = 7
	case AddrIPv6:
		n = 19
	case AddrDomain:
		n = 4 + len(addr.Host)
	default:
		n = 7
	}
	return
}

func (addr *Addr) String() string {
	return net.JoinHostPort(addr.Host, strconv.Itoa(int(addr.Port)))
}

//
// class def Addr end
//

//
// class Reply start
//
const (
	Succeeded uint8 = iota
	Failure
	NotAllowed
	NetUnreachable
	HostUnreachable
	ConnRefused
	TTLExpired
	CmdUnsupported
	AddrUnsupported
)

/*
The SOCKSv5 reply
 +----+-----+-------+------+----------+----------+
 |VER | REP |  RSV  | ATYP | BND.ADDR | BND.PORT |
 +----+-----+-------+------+----------+----------+
 | 1  |  1  | X'00' |  1   | Variable |    2     |
 +----+-----+-------+------+----------+----------+
*/
type Reply struct {
	Rep  uint8
	Addr *Addr
}

func NewReply(rep uint8, addr *Addr) *Reply {
	return &Reply{
		Rep:  rep,
		Addr: addr,
	}
}

func ReadReply(r io.Reader) (*Reply, error) {
	var b [262]byte

	n, err := io.ReadAtLeast(r, b[:], 5)
	if err != nil {
		return nil, err
	}

	if b[0] != 5 {
		return nil, ErrBadVersion
	}

	reply := &Reply{
		Rep: b[1],
	}

	atype := b[3]
	length := 0
	switch atype {
	case AddrIPv4:
		length = 10
	case AddrIPv6:
		length = 22
	case AddrDomain:
		length = 7 + int(b[4])
	default:
		return nil, ErrBadAddrType
	}

	if n < length {
		if _, err := io.ReadFull(r, b[n:length]); err != nil {
			return nil, err
		}
	}

	addr := new(Addr)
	if err := addr.Decode(b[3:length]); err != nil {
		return nil, err
	}
	reply.Addr = addr

	return reply, nil
}

func (r *Reply) Write(w io.Writer) (err error) {
	var b [262]byte

	b[0] = 5
	b[1] = r.Rep
	b[2] = 0        //rsv
	b[3] = AddrIPv4 // default
	length := 10
	b[4], b[5], b[6], b[7], b[8], b[9] = 0, 0, 0, 0, 0, 0 // reset address field

	if r.Addr != nil {
		n, _ := r.Addr.Encode(b[3:])
		length = 3 + n
	}
	_, err = w.Write(b[:length])

	return
}

func (r *Reply) String() string {
	addr := r.Addr
	if addr == nil {
		addr = &Addr{}
	}
	return fmt.Sprintf("5 %d 0 %d %s",
		r.Rep, addr.Type, addr.String())
}

func main() {
	fmt.Println("这是一个简单的socks服务器,监听在端口10080")
	fmt.Println("Listening on 0.0.0.0:10080")
	listener, err := net.Listen("tcp", "0.0.0.0:10080")
	if err != nil {
		fmt.Println("错误，无法启动监听")
		return
	}
	go func() {
		for {
			conn, err := listener.Accept()
			if err != nil {
				break
			}
			saddr := conn.RemoteAddr().String()
			fmt.Println("connection setup, from:" + saddr)
			host, _, err := net.SplitHostPort(saddr)
			if err != nil {
				fmt.Println("无法从中解析地址端口:" + saddr)
				continue
			}

			if !isInWhiteList(host) {
				fmt.Printf("拒绝服务%s，不在白名单中。\n", host)
				conn.Close()
				continue
			}

			go DoClient(conn)
		}
	}()

	runHttpServer()
	select {}
}

func isInWhiteList(host string) bool {
	for i := 0; i < len(gUsers); i++ {
		for k, v := range gUsers[i].Expires {
			if v.IP == host {
				if time.Now().After(v.ExpireTime) {
					fmt.Println("在白名单中但是超时了,删除它。")
					gUsers[i].Expires = append(gUsers[i].Expires[:k], gUsers[i].Expires[k+1:]...)
					return false
				} else {
					return true
				}
			}
		}
	}
	return false
}

func DoClient(conn net.Conn) {
	defer conn.Close()
	var buf [256]byte
	/*
		* 版本标识符/方法选择消息（version identifier/method selection message）
		* 对于此版本的协议，VER字段被固定设置为 X'05'。 NMETHODS字段表示在METHODS字段中的方法标识符（以字节为单位）的数量。
		* 当前METHOD定义的值是：
		* X'00' 无需认证（NO AUTHENTICATION REQUIRED）
		* X'01' GSSAPI
		* X'02' 用户名/密码（USERNAME/PASSWORD）
		* X'03' 至 X'7F' IANA 分配（IANA ASSIGNED）
		* X'80' 到 X'FE' 保留用于私人方法（RESERVED FOR PRIVATE METHODS）
		* X'FF' 没有可接受的方法（NO ACCEPTABLE METHODS）
		+----+----------+----------+
		|VER | NMETHODS | METHODS  |
		+----+----------+----------+
		| 1  | 1        | 1 to 255 |
		+----+----------+----------+
	*/
	cnt, err := io.ReadAtLeast(conn, buf[:], 3)
	if err != nil {
		fmt.Println("METHOD SELECT: 无效的socks5请求!")
		return
	}
	if buf[0] != 5 {
		fmt.Println("METHOD SELECT: 无效的版本标识")
		return
	}
	methodBytes := buf[1]
	msgLength := 1 + 1 + int(methodBytes)
	if cnt < msgLength {
		_, err := io.ReadFull(conn, buf[cnt:msgLength])
		if err != nil {
			fmt.Println("METHOD SELECT: 无法读取完整的数据头")
			return
		}
	}
	// 发送选择方法：这个版本只返回无账号密码验证模式
	selectReply := [2]byte{5, 0}
	n, err := conn.Write(selectReply[:])
	if err != nil {
		fmt.Println("METHOD SELECT: 发送回复失败了（1）。")
		return
	}
	if n <= 0 {
		fmt.Println("METHOD SELECT: 发送回复失败了（1）。")
		return
	}

	//
	// 未来在这里添加认证步骤
	//
	/*
		Socks5 REQUEST.
		+----+-----+-------+------+----------+----------+
		|VER | CMD | RSV   | ATYP | DST.ADDR | DST.PORT |
		+----+-----+-------+------+----------+----------+
		| 1  | 1   | X'00' | 1    | Variable | 2        |
		+----+-----+-------+------+----------+----------+
		其中：
		VER 协议版本号: X'05'
		CMD
			CONNECT X'01'
			BIND X'02'
			UDP ASSOCIATE X'03'
		RSV 保留未用（RESERVED）
		ATYP 后续地址（即DST.ADDR）的地址类型（address type of following address）
			IP V4 address: X'01'
			DOMAINNAME: X'03'
			IP V6 address: X'04'
		DST.ADDR 期望的目标地址
		DST.PORT 期望的目标端口，网络字节序
		SOCKS服务器通常会根据源地址和目标地址评估请求，并根据请求类型返回一条或多条回复消息。

		Reply:
		+----+-----+-------+------+----------+----------+
		|VER | REP | RSV   | ATYP | BND.ADDR | BND.PORT |
		+----+-----+-------+------+----------+----------+
		| 1  | 1   | X'00' | 1    | Variable | 2        |
		+----+-----+-------+------+----------+----------+
		VER 协议版本：X'05'
		REP 回复字段：
			X'00' 成功
			X'01' 普通SOCKS服务器故障
			X'02' 连接不被规则集允许
			X'03' 网络不可达
			X'04' 主机不可达
			X'05' 连接被拒绝
			X'06' TTL 过期
			X'07' 不支持的命令
			X'08' 不支持的地址类型
			X'09' 到 X'FF' 未分配
		RSV 保留未用，总是为0
		ATYP 后续地址(即BND.ADDR)的地址类型
			IP V4 地址：X'01'
			域名：X'03'
			IP V6 地址：X'04'
		BND.ADDR 服务器绑定的地址
		BND.PORT 服务器绑定的端口（网络字节序
	*/
	cnt, err = io.ReadAtLeast(conn, buf[:], 5)
	if err != nil {
		fmt.Println("REQUEST:无效的请求!")
		return
	}

	if buf[0] != 5 {
		fmt.Println("REQUEST:无效的版本标识")
		return
	}

	cmd := buf[1]
	atype := buf[3]
	length := 0
	switch atype {
	case AddrIPv4:
		length = 10
	case AddrIPv6:
		length = 22
	case AddrDomain:
		length = 7 + int(buf[4])
	default:
		fmt.Println("REQUEST:无效的地址标识")
		return
	}

	if cnt < length {
		_, err := io.ReadFull(conn, buf[cnt:length])
		if err != nil {
			fmt.Println("REQUEST:无法读取完整的数据头")
			return
		}
	}

	addr := new(Addr)
	if err := addr.Decode(buf[3:length]); err != nil {
		fmt.Println("REQUEST:读取目标地址失败")
		return
	}

	// 这里的address可能是IP地址也可能是一个域名，这取决于socks5客户端的实现。
	// 对于google chrome通常传递域名，对于python requests模块则通常传递IP地址。
	strAddress := addr.String()
	fmt.Println("目标地址:" + strAddress)

	// switch cmd {
	// case CmdConnect:
	// 	fmt.Println("CMD:收到一个CONN请求")
	// case CmdBind:
	// case CmdMuxBind:
	// case CmdUdp:
	// case CmdUDPTun:
	// default:
	// 	fmt.Println("CMD:暂不支持该CMD类型")
	// 	resp := NewReply(CmdUnsupported, nil)
	// 	resp.Write(conn)
	// 	return
	// }
	if cmd != CmdConnect {
		fmt.Println("REQUEST::暂不支持该CMD类型")
		resp := NewReply(CmdUnsupported, nil)
		resp.Write(conn)
		return
	}

	//
	// 拿到客户端想连结的服务器地址后，尝试去连接
	//
	ipAddr, err := net.ResolveIPAddr("ip", addr.Host)
	if err != nil {
		fmt.Println("REQUEST:无法解析连结远端服务器")
		return
	}
	remote := net.TCPAddr{IP: ipAddr.IP, Port: int(addr.Port)}
	connRemote, err := net.DialTCP("tcp", nil, &remote)
	if err != nil {
		fmt.Printf("连接远端服务器%s失败了，错误:%s", ipAddr.String(), err.Error())
		return
	}
	fmt.Println("连接服务器成功!" + connRemote.RemoteAddr().String())
	defer connRemote.Close()

	// 告知socks5客户端，已成功建立代理通道。
	resp := NewReply(Succeeded, nil)
	err = resp.Write(conn)
	if err != nil {
		fmt.Println("向客户端发送回应失败了:" + err.Error())
		return
	}

	// 创建两个go routine，其中一个将读取客户端发来的数据并发往远端服务器，
	// 另一个则将读取远端服务器发来的数据并发往客户端。
	errc := make(chan error, 1)
	go func() {
		// 如果对默认的内存大小或自动释放机制不满意，自己实现一个bufferPool，然后用io.CopyBuffer(conn, connRemote, pooledBuffer)
		_, errcode := io.Copy(connRemote, conn)
		errc <- errcode
	}()

	go func() {
		_, errcode := io.Copy(conn, connRemote)
		errc <- errcode
	}()

	err = <-errc
	if err != nil && err != io.EOF {
		fmt.Printf("连结意外断开了，错误:%s", err.Error())
	}
	return
}
